#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>

#include <pcl/io/pcd_io.h>
#include <pcl/io/ply_io.h>
#include <pcl/point_types.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <Eigen/Dense>

// nlohmann json
#include <json.hpp>
using json = nlohmann::json;

void readKittiBin(const std::string& bin_path, pcl::PointCloud<pcl::PointXYZI>::Ptr cloud) {
    std::ifstream input(bin_path.c_str(), std::ios::binary);
    if (!input.good()) {
        std::cerr << "Error opening file: " << bin_path << std::endl;
        return;
    }
    input.seekg(0, std::ios::beg);
    while (input.peek() != EOF) {
        float x, y, z, intensity;
        input.read((char*)&x, sizeof(float));
        input.read((char*)&y, sizeof(float));
        input.read((char*)&z, sizeof(float));
        input.read((char*)&intensity, sizeof(float));
        pcl::PointXYZI p;
        p.x = x; p.y = y; p.z = z; p.intensity = intensity;
        cloud->push_back(p);
    }
    input.close();
}

// 保持原有
void draw_kitti_box(pcl::visualization::PCLVisualizer::Ptr viewer,
                 const std::vector<float>& box,
                 std::string box_name = "bbox",
                 double r=1, double g=0, double b=0)
{
    float cx = box[0], cy = box[1], cz = box[2];
    float l = box[3], w = box[4], h = box[5], yaw = box[6];
    // ... 保持draw_kitti_box内容不变
    float x_corners[8] = { l/2,  l/2, -l/2, -l/2,  l/2,  l/2, -l/2, -l/2};
    float y_corners[8] = { w/2, -w/2, -w/2,  w/2,  w/2, -w/2, -w/2,  w/2};
    float z_corners[8] = {   0,    0,    0,    0,    h,    h,    h,    h};

    Eigen::Matrix3f R;
    R = Eigen::AngleAxisf(yaw, Eigen::Vector3f::UnitZ());
    std::vector<Eigen::Vector3f> corners_world(8);

    for (int i = 0; i < 8; ++i) {
        Eigen::Vector3f pt_local(x_corners[i], y_corners[i], z_corners[i]);
        Eigen::Vector3f pt_world = R * pt_local + Eigen::Vector3f(cx, cy, cz);
        corners_world[i] = pt_world;
    }

    std::vector<std::pair<int, int>> lines = {
        {0,1}, {1,2}, {2,3}, {3,0}, // bottom
        {4,5}, {5,6}, {6,7}, {7,4}, // top
        {0,4}, {1,5}, {2,6}, {3,7}  // vertical
    };

    for (size_t i = 0; i < lines.size(); ++i) {
        std::stringstream ss;
        ss << box_name << "_edge_" << i;
        viewer->addLine(
            pcl::PointXYZ(
                corners_world[lines[i].first](0),
                corners_world[lines[i].first](1),
                corners_world[lines[i].first](2)),
            pcl::PointXYZ(
                corners_world[lines[i].second](0),
                corners_world[lines[i].second](1),
                corners_world[lines[i].second](2)),
            r, g, b, ss.str()
        );
    }
}

void read_bboxes3d_scores_from_json(
    const std::string& json_path,
    std::vector<std::vector<float>>& out_bboxes,
    std::vector<float>& out_scores)
{
    out_bboxes.clear(); out_scores.clear();
    std::ifstream jfile(json_path);
    if (!jfile.is_open()){
        std::cerr << "Cannot open json file: " << json_path << std::endl;
        return;
    }
    json jdata;
    jfile >> jdata;
    jfile.close();
    if (!jdata.contains("bboxes_3d") || !jdata.contains("scores_3d")) {
        std::cerr << "JSON missing keys!" << std::endl;
        return;
    }
    // 读取
    for (auto& x : jdata["bboxes_3d"])
        out_bboxes.push_back(x.get<std::vector<float>>());
    for (auto& s : jdata["scores_3d"])
        out_scores.push_back((float)s);
}

int main()
{
    // 1. 读取点云
    pcl::PointCloud<pcl::PointXYZI>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZI>());
    readKittiBin("/home/chenhao/kitti3d_segmentation_benchmark/data/bin/000000.bin", cloud);

    // 2. 读取json
    std::string json_path = "/home/chenhao/kitti3d_segmentation_benchmark/output/000000.json";
    std::vector<std::vector<float>> bboxes_3d;
    std::vector<float> scores_3d;
    read_bboxes3d_scores_from_json(json_path, bboxes_3d, scores_3d);

    // 3. 根据score过滤
    std::vector<std::vector<float>> filtered_bboxes;
    for(size_t i=0; i<bboxes_3d.size(); ++i){
        if(scores_3d[i] > 0.5){
            filtered_bboxes.push_back(bboxes_3d[i]);
            // 打印框
            std::cout << "Box " << i << " (score: " << scores_3d[i] << "): [ ";
            for(auto v: bboxes_3d[i]) std::cout << v << " ";
            std::cout << "]\n";
        }
    }

    // 4. 可视化
    pcl::visualization::PCLVisualizer::Ptr viewer(new pcl::visualization::PCLVisualizer("PointCloud with BBoxes"));
    viewer->setBackgroundColor(0, 0, 0);
    viewer->addPointCloud<pcl::PointXYZI>(cloud, "cloud");
    viewer->setPointCloudRenderingProperties(pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 1, "cloud");
    for (size_t i = 0; i < filtered_bboxes.size(); ++i) {
        draw_kitti_box(viewer, filtered_bboxes[i], "bbox_" + std::to_string(i), 1, 0, 0);
    }
    while (!viewer->wasStopped())
        viewer->spinOnce(100);
    return 0;
}